import json

from app.provider.knowledge_provider import KnowledgeProvider
from zhanshop.hash import Hash
from zhanshop.helper.str import Str
from zhanshop.json import Json


class KnowledgeService():
    models = None
    def __init__(self):
        self.models = {}

    def getModel(self, database, table, modelParam):
        """
        获取模型
        :param database:
        :param table:
        :param modelParam:
        :return:
        """
        key = Hash.md5(Json.encode(modelParam))
        model = self.models.get(key, None)
        #if (model == None):
        model = KnowledgeProvider(database, table, modelParam)
        self.models[key] = model
        return model

    def getRepository(self, database: str, table: str, question: str, limit: int = 100, modelParam: json = None):
        """
        查询知识库
        :param database:
        :param table:
        :param question:
        :param modelParam:
        :return:
        """
        if(modelParam == None or modelParam == ""): modelParam = {
            "SentenceTransformer": {"model": "shibing624/text2vec-base-chinese"},
            "chromadb": {"metadata": {"hnsw:space": "cosine"}},
            "CrossEncoder": {"model_name_or_path": "cross-encoder/mmarco-mMiniLMv2-L12-H384-v1"}}
        model = self.getModel(database, table, modelParam)
        searchResults = model.query(question, limit)
        #return searchResults
        return model.getScores(question, searchResults)

    def postRepository(self, database: str, table: str, documents: list, metadatas: list, modelParam: json = None):
        """
        添加知识库
        :param database:
        :param table:
        :param documents:
        :param metadatas:
        :param modelParam:
        :return:
        """
        if (modelParam == None or modelParam == ""): modelParam = {
            "SentenceTransformer": {"model": "shibing624/text2vec-base-chinese"},
            "chromadb": {"metadata": {"hnsw:space": "cosine"}},
            "CrossEncoder": {"model_name_or_path": "cross-encoder/mmarco-mMiniLMv2-L12-H384-v1"}}
        model = self.getModel(database, table, modelParam)
        documents = list(filter(lambda x: x != '', documents))
        model.batchAdd(documents, metadatas)
        return {}